
#include "arg_max_with_value_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <cmath>
#include <stdio.h>


namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    ArgMaxWithValueTilingData tiling;

    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();
    uint32_t dims = context->GetInputShape(0)->GetStorageShape().GetDimNum();
    const int64_t *axesVal = context->GetAttrs()->GetInt(0);
    const bool *keep_dims = context->GetAttrs()->GetBool(1);
    int64_t axes = *axesVal;
    //确定操作步长
    uint32_t stride = 1;
    uint32_t shapeDim[12];
    if(axes<0) axes += dims;
    for(int i=0;i<dims;i++) shapeDim[i] = context->GetInputShape(0)->GetStorageShape().GetDim(i);
    for(int i=axes+1;i<dims;i++) stride *= shapeDim[i];
    //确定操作Iter数
    uint32_t iterStep = stride * shapeDim[axes];
    uint32_t nIter = 1;
    for(int i=0;i<axes;i++) nIter *= shapeDim[i];
    uint32_t resLength = stride * nIter;
    // printf("dims: %d axes:%d stride %d nIter %d\n",dims, axes, stride, nIter);
    tiling.set_iterStep(iterStep);          //组间地址偏移
    tiling.set_nIter(nIter);                //组数
    tiling.set_totalLength(totalLength);    //元素数
    tiling.set_resLength(resLength);        //dst Tensor大小
    tiling.set_stride(stride);              //比较步长(也是单组内需要比较的次数)
    tiling.set_dims(dims);                  
    tiling.set_axes(axes);
    tiling.set_keep_dims(*keep_dims);
    tiling.set_axesDim(shapeDim[axes]);     //每次需要比较的数值个数
    // printf("Tiling done.\n: totalLength=%u\n resLength=%u\n dims=%u\n axes=%ld\n stride=%u\n nIter=%u\n IterStep=%u\n axesDim=%u\n", \
    totalLength, resLength, dims, axes, stride, nIter, iterStep,shapeDim[axes]);
    const gert::Shape vec = context->GetInputShape(0)->GetOriginShape();
    auto aicNum = ascendcPlatform.GetCoreNumAic();
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    aivNum = 32;
    // printf("aicNum=%d aivNum=%d\n", aicNum, aivNum);
    int32_t smallDataNum, bigDataNum;
    int32_t tailBlockNum = 0;
    
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType());

    // 处理half
    // printf("axes: %d\n", axes);
    context->SetTilingKey(0);
    if (axes == 0 && inputBytes==1) {
        if (aivNum >= stride) {
            aivNum = stride;
            smallDataNum = 1;
            bigDataNum = 1;
        } else {
            // printf("%d %d = %d\n",stride,aivNum,smallDataNum);
            smallDataNum = stride/aivNum;
            bigDataNum = smallDataNum + 1;
            tailBlockNum = stride%aivNum; // 需要处理更大块数据
        }
    } else {
        if (context->GetInputDesc(0)->GetDataType() == 0 && axes == dims - 1) {
            if (shapeDim[axes] >= 500) {
                context->SetTilingKey(1);
            }
            if (aivNum >= resLength) {
                aivNum = resLength;
                smallDataNum = 1;
                bigDataNum = 1;
          } else{
                smallDataNum = resLength/aivNum;
                bigDataNum = smallDataNum + 1;
                tailBlockNum = resLength%aivNum;
                tiling.set_magicNum(1000*resLength + shapeDim[axes] );
                // context->SetTilingKey(1);
            }
        } else {
            // context->SetTilingKey(2);
            if (aivNum >= resLength) {
                aivNum = resLength;
                smallDataNum = 1;
                bigDataNum = 1;
            } else {
                // 根据完整的前后进行划分
                smallDataNum = resLength/aivNum;
                bigDataNum = smallDataNum + 1;
                tailBlockNum = resLength%aivNum; // 需要处理更大块数据
            }
        }
    }
    // printf("small:%d big:%d resLength:%d\n", smallDataNum, bigDataNum, resLength);
    // printf("ubSize=%u\n", ub_size);
    // printf("small :%d %d\n", smallDataNum, bigDataNum);
    // 设置使用全部的vector core
    tiling.set_smallDataNum(smallDataNum);
    tiling.set_bigDataNum(bigDataNum);     //每次需要比较的数值个数
    tiling.set_tailBlockNum(tailBlockNum);
    context->SetBlockDim(aivNum);  
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ArgMaxWithValue : public OpDef {
public:
    explicit ArgMaxWithValue(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("indice")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("values")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("dimension").Int();
        this->Attr("keep_dims").AttrType(OPTIONAL).Bool(false);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(ArgMaxWithValue);
}
